{
  "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reaching multiple targets with a manipulator\n",
    "The objective of this exercise is to reach multiple targets with a manipulator.\n",
    "\n",
    "We provide a basic example for reaching one point, and you have to modify it for sequence of multiple targets. Below it is the basic example, there we'll guide you to the final result.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %load arm_example.py\n",
    "import crocoddyl\n",
    "import pinocchio\n",
    "import numpy as np\n",
    "import example_robot_data\n",
    "\n",
    "robot = example_robot_data.load(\"talos_arm\")\n",
    "robot_model = robot.model\n",
    "\n",
    "DT = 1e-3\n",
    "T = 25\n",
    "target = np.array([0.4, 0.0, 0.4])\n",
    "\n",
    "display = crocoddyl.GepettoDisplay(robot)\n",
    "display.robot.viewer.gui.addSphere(\n",
    "    \"world/point\", 0.05, [1.0, 0.0, 0.0, 1.0]\n",
    ")  # radius = .1, RGBA=1001\n",
    "display.robot.viewer.gui.applyConfiguration(\n",
    "    \"world/point\", target.tolist() + [0.0, 0.0, 0.0, 1.0]\n",
    ")  # xyz+quaternion\n",
    "display.robot.viewer.gui.refresh()\n",
    "\n",
    "# Create the cost functions\n",
    "state = crocoddyl.StateMultibody(robot.model)\n",
    "goalTrackingCost = crocoddyl.CostModelResidual(\n",
    "    state,\n",
    "    crocoddyl.ResidualModelFrameTranslation(\n",
    "        state, robot_model.getFrameId(\"gripper_left_joint\"), target\n",
    "    ),\n",
    ")\n",
    "xRegCost = crocoddyl.CostModelResidual(state, crocoddyl.ResidualModelState(state))\n",
    "uRegCost = crocoddyl.CostModelResidual(state, crocoddyl.ResidualModelControl(state))\n",
    "\n",
    "# Create cost model per each action model\n",
    "runningCostModel = crocoddyl.CostModelSum(state)\n",
    "terminalCostModel = crocoddyl.CostModelSum(state)\n",
    "\n",
    "# Then let's added the running and terminal cost functions\n",
    "runningCostModel.addCost(\"gripperPose\", goalTrackingCost, 1e2)\n",
    "runningCostModel.addCost(\"stateReg\", xRegCost, 1e-4)\n",
    "runningCostModel.addCost(\"ctrlReg\", uRegCost, 1e-7)\n",
    "terminalCostModel.addCost(\"gripperPose\", goalTrackingCost, 1e5)\n",
    "terminalCostModel.addCost(\"stateReg\", xRegCost, 1e-4)\n",
    "terminalCostModel.addCost(\"ctrlReg\", uRegCost, 1e-7)\n",
    "\n",
    "# Create the actuation model\n",
    "actuationModel = crocoddyl.ActuationModelFull(state)\n",
    "\n",
    "# Create the action model\n",
    "runningModel = crocoddyl.IntegratedActionModelEuler(\n",
    "    crocoddyl.DifferentialActionModelFreeFwdDynamics(\n",
    "        state, actuationModel, runningCostModel\n",
    "    ),\n",
    "    DT,\n",
    ")\n",
    "terminalModel = crocoddyl.IntegratedActionModelEuler(\n",
    "    crocoddyl.DifferentialActionModelFreeFwdDynamics(\n",
    "        state, actuationModel, terminalCostModel\n",
    "    )\n",
    ")\n",
    "# runningModel.differential.armature = 0.2 * np.ones(state.nv)\n",
    "# terminalModel.differential.armature = 0.2 * np.ones(state.nv)\n",
    "\n",
    "# Create the problem\n",
    "q0 = np.array([2.0, 1.5, -2.0, 0.0, 0.0, 0.0, 0.0])\n",
    "x0 = np.concatenate([q0, pinocchio.utils.zero(state.nv)])\n",
    "problem = crocoddyl.ShootingProblem(x0, [runningModel] * T, terminalModel)\n",
    "\n",
    "# Creating the DDP solver for this OC problem, defining a logger\n",
    "ddp = crocoddyl.SolverDDP(problem)\n",
    "ddp.setCallbacks([crocoddyl.CallbackVerbose()])\n",
    "\n",
    "# Solving it with the DDP algorithm\n",
    "ddp.solve()\n",
    "\n",
    "# Visualizing the solution in gepetto-viewer\n",
    "display.displayFromSolver(ddp)\n",
    "\n",
    "robot_data = robot_model.createData()\n",
    "xT = ddp.xs[-1]\n",
    "pinocchio.forwardKinematics(robot_model, robot_data, xT[: state.nq])\n",
    "pinocchio.updateFramePlacements(robot_model, robot_data)\n",
    "print(\n",
    "    \"Finally reached = \",\n",
    "    robot_data.oMf[robot_model.getFrameId(\"gripper_left_joint\")].translation.T,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## I. DifferentialActionModel for Pinocchio ABA\n",
    "This scenario uses an action model that computes 2nd order differential dynamics with Pinocchio. Note that it can accept several cost models. This action model is tailored for robot applications, and at the same time, it's modular since:\n",
    " - you can modify the robot dynamics by changing Pinocchio model, and\n",
    " - you can formulate any cost function by simply adding running a terminal costs.\n",
    "\n",
    "## II. Cost models\n",
    "\n",
    "A cost model computes a scalar cost value and its gradient and Hessian. All the models implemented are computing a cost residual and are computing the Hessian with the Gauss approximation.\n",
    "\n",
    "We implemented reusable cost models for controlling \n",
    " - a frame placement (translation or velocity),\n",
    " - the center of mass position, and \n",
    " - state  and control spaces.\n",
    "\n",
    "In the example above, we used the CostModelFrameTranslation which defines a 3d position task, and the state and control regularizers.\n",
    "\n",
    "As for any cost model in crocoddyl, if you write your own cost model you need to create a data class for your cost function. The cost data must be created from a pinocchio data (the rational is that the pinocchio data used to compute the dynamics should be re-used to compute the cost).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataCollector = crocoddyl.DataCollectorMultibody(robot.data)\n",
    "trackData = goalTrackingCost.createData(dataCollector)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### II.a Frame position cost\n",
    "\n",
    "You define a frame ID and the reference position as a 3D array. The cost is the distance between the frame and the target. This cost depends on $\\mathbf{x}$ (specifically the configuration $\\mathbf{q}$). You can double check the 0s in its gradient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pinocchio.updateFramePlacements(robot.model, robot.data)\n",
    "pinocchio.computeJointJacobians(robot.model, robot.data, xT[: state.nq])\n",
    "goalTrackingCost.calc(trackData, x0)\n",
    "goalTrackingCost.calcDiff(trackData, x0)\n",
    "print(trackData.Lx, trackData.Lu)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### II.b State cost\n",
    "In this part of the tutorial you must define a State model. It defines \n",
    " - the dimension of the state and its tangent, and\n",
    " - the exponential/integrate and difference/log operators.\n",
    "The operators can described using Pinocchio functions. And the exercite consists on adding them into your State class. Please note crocoddyl has abstract functions for this.\n",
    "\n",
    "The state cost uses a reference in state space (State.zero() by default). The cost is the distance, computed with state.difference between the current state and the reference. Hence, with this cost, we regularize both position and velocity.\n",
    "\n",
    "### II.c Control cost\n",
    "\n",
    "The control cost uses a control reference as in the state cost. The cost is the distance the current control and the reference. Hence the cost regularizes torque commands."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### II.d Add cost models to the differential action model\n",
    "Each time we want to include a new cost function, we use addCost function inside our DAM. In this function you're also able its weight."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## III. Create the problem with integrated action model\n",
    "Differential action models describe cost and dynamics in continuous-time, however our optimal control solvers work in discrete-time. We have created the integrated action model in order to deal with this.\n",
    "\n",
    "In the previous code, we have used an abstract class that uses simpletic Euler rules. In the cartpole exercise you have learnt how to use integrated action models for your problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## IV. Callbacks\n",
    "\n",
    "Callback functions are needed for analysing and debugging the performance of the solver for your specific problem.\n",
    "For problems defined with Pinocchio, you can display the robot trajectory per each iterate by including CallbackDisplay. With this callback, you can display robot motions with different rates. Additionally, CallbackVerbose prints a message that allows us to understand the behaviour of the solver.\n",
    "\n",
    "Generally speaking, an user is able to describe any callback function. This function will be run once per iterate and it has access to all data."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## VI. Modifying the example\n",
    "\n",
    "Start by defining several targets (let's say 4 targets, all at x=0.4, and at y and z being either 0 or 0.4), and display then in the viewer.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The shooting problem will be composed of 4 sequences of action models. Each sequence consists on T shooting \"running\" nodes and 1 terminal node. The running nodes mostly have regularization terms, while the terminal nodes have a strong cost toward the respective target.\n",
    "\n",
    "[ R1,R1,R1 ... R1,T1, R2,R2 .... R2, T2, R3 ... R3, T3, R4 ... R4 ] , T4\n",
    "\n",
    "First create 4 running models and 4 terminal models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then you need to add a position cost, and state and control regularization to each running action model. Please  note that for terminal action model is only needed the position cost. Additionally, in the running models, the position cost should be low, and it should be high in the terminal models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now create a shooting problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "seq0 = [runningModels[0]] * T + [terminalModels[0]]\n",
    "seq1 = [runningModels[1]] * T + [terminalModels[1]]\n",
    "seq2 = [runningModels[2]] * T + [terminalModels[2]]\n",
    "seq3 = [runningModels[3]] * T\n",
    "problem = crocoddyl.ShootingProblem(x0, seq0 + seq1 + seq2 + seq3, terminalmodel[3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a DDP solver for this problem and run it. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ddp = crocoddyl.SolverDDP(problem)\n",
    "ddp.solve()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Well, it should not work, at least no on the first shot. The DDP solver is likely not strong enough to accept the random weights that you have selected. \n",
    "\n",
    "If it is working nicely from the first shot, display it in the viewer and go take a coffee. But you will likely have to tweak the gains to make it work.\n",
    "\n",
    "**It is suggested to first optimize only sequence 1. When you are happy with it, add sequence 2 and optimize again, etc.**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## V. Penalty\n",
    "The solver works with double precisions, so it is quite robust to high weight. 10000 is likely to be accepted for example. But if you make the problem too difficult, the solver will break. \n",
    "\n",
    "In that case, you can implement a simple penalty solver by setting the weight to be 10**i, and creating a for loop to explore i from 0 to 5. At each iteration of the loop, run the solver from the previous solution and for few iterations only."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(1, 6):\n",
    "    for m in terminalModels:\n",
    "        m.costs.costs[\"gripperPose\"].weight = 10**i\n",
    "    ddp.solve(ddp.xs, ddp.us, 10)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "3.10.4 (v3.10.4:9d38120e33, Mar 23 2022, 17:29:05) [Clang 13.0.0 (clang-1300.0.29.30)]"
  },
  "vscode": {
   "interpreter": {
    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
